package com.zy;

import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.ClassPathDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByCharacterSplitter;
import dev.langchain4j.data.document.splitter.DocumentByRegexSplitter;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;

import java.util.List;

/**
 * @program: AI_langchain4j
 * @description: 从某个位置加载文档, 向量化，存向量库，检索
 * @author: zy
 * @create: 2025-07-05 16:43
 */
public class RAG_App3 {
    public static void main(String[] args) {
        //1. 文件加载器有多种实现，这里使用ClassPathDocumentLoader,还有FileDocumentLoader,UrlDocumentLoader等
        //DocumentParser解析器有多种实现，这里使用TextDocumentParser,还有MarkdownDocumentParser,HtmlDocumentParser等
        Document doc= ClassPathDocumentLoader.loadDocument("rag/a.txt", new TextDocumentParser( ));
        //System.out.println(doc.text());
        //2. 分词器: 决定了文本如何被分割成更小的单元，这些单元可以是单词、句子、段落或其他更细粒度的文本片段。
        // DocumentSplitter文档拆分器有多种实现，这里使用RecursiveCharacterTextSplitter,还有TokenTextSplitter等
//        DocumentByCharacterSplitter splitter=new DocumentByCharacterSplitter(
//                100,   //每段最长字数
//                10
//        );
        DocumentByRegexSplitter splitter=new DocumentByRegexSplitter(
                "\\n\\d+\\.\\s+",
                "\\n",     //保留换行符作为段落连接符
                90,
                1,
                new DocumentByCharacterSplitter(100,10)
        );
        List<TextSegment> segments=splitter.split( doc);
        //System.out.println(   segments );

        //3. 向量化器: 负责将文本片段转换为数值向量，这些向量可以被机器学习模型用于文本分析、检索和生成。
        // 向量化器有多种实现，这里使用OpenAIEmbeddingService,还有HuggingFaceEmbeddingService等
        QwenEmbeddingModel qwenEmbeddingModel = QwenEmbeddingModel.builder()
                .apiKey(System.getenv("ALI_AI_KEY"))
                .build();
        //向量化
        List<Embedding> embeddings=qwenEmbeddingModel.embedAll( segments ).content();

        //4. 保存到向量数据库中 此处还是用 InMemoryEmbeddingStore ，实际应用中可以使用其他数据库
        // 向量数据库有多种实现，这里使用InMemoryEmbeddingStore,还有ElasticsearchEmbeddingStore等
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore();
        embeddingStore.addAll( embeddings,segments );



        /////////////////////////////阶段二:检索.


        //5. 对查询关键词进行向量化
        Response<Embedding> embed=qwenEmbeddingModel.embed("装卸工工作职责");
        //创建请求
        EmbeddingSearchRequest build=EmbeddingSearchRequest.builder()
                .queryEmbedding(embed.content())
                .maxResults( 2 )      //返回最相似的一条数据
                .build();
        //6. 查询
        EmbeddingSearchResult<TextSegment> segmentEmbeddingSearchResult=embeddingStore.search( build );
        segmentEmbeddingSearchResult.matches().forEach(   embeddingMatch->{
            System.out.println("========================");
            System.out.println(   embeddingMatch.score());
            System.out.println( embeddingMatch.embedded().text() );
            System.out.println("========================");
        });


        //////////////////////////阶段三: 检索增强阶段，用户发问->问题向量化->向量数据库查询->返回上下文->prompt->模型生成答案->返回答案
        //大语言模型 ,
        ChatModel chatModel= QwenChatModel.builder()
                .apiKey(System.getenv("ALI_AI_KEY"))
                .modelName("qwen-turbo")
                .build();
        //内容检索器
        ContentRetriever contentRetriever= EmbeddingStoreContentRetriever.builder()
                .embeddingStore( embeddingStore )
                .embeddingModel( qwenEmbeddingModel )
                .maxResults(5)
                .minScore(0.7)
                .build();
        //为Assistant生成动态代理对象 chat -->对话内容存储 ChatMemory -->聊天
        Assistant assistant= AiServices.builder(   Assistant.class)
                .chatModel(  chatModel    )
                .contentRetriever( contentRetriever )   ///   内容检索
                .build();
        System.out.println("*****************************");
        System.out.println(   assistant.chat( "装卸工工作职责" ) );

    }

    public interface Assistant{
        String chat(String userMessage);

    }

}
